package com.theme.passport.aspect;

import com.theme.common.mvc.view.R;
import net.jodah.expiringmap.ExpirationPolicy;
import net.jodah.expiringmap.ExpiringMap;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * 接口限流切面
 * @author cy
 * @date 2023/2/24 14:25
 */
@Aspect
@Component
public class CurrentLimitRequestAspect {

    private static ConcurrentHashMap<String, ExpiringMap<String, Integer>> concurrentHashMap = new ConcurrentHashMap<>();

    // 定义切点
    // 让所有有@CurrentLimitRequest注解的方法都执行切面方法
    @Pointcut("@annotation(currentLimitRequest)")
    public void excludeService(CurrentLimitRequest currentLimitRequest) {
    }

    @Around("excludeService(currentLimitRequest)")
    public Object doAround(ProceedingJoinPoint point, CurrentLimitRequest currentLimitRequest) throws Throwable {
        // 获得request对象
        RequestAttributes ra = RequestContextHolder.getRequestAttributes();
        ServletRequestAttributes sra = (ServletRequestAttributes) ra;
        assert sra != null;
        HttpServletRequest request = sra.getRequest();

        // 获取Map对象， 如果没有则返回默认值
        // 第一个参数是key， 第二个参数是默认值
        ExpiringMap<String, Integer> map = concurrentHashMap.getOrDefault(request.getRequestURI(), ExpiringMap.builder().variableExpiration().build());
        Integer uCount = map.getOrDefault(request.getRemoteAddr(), 0);

        if (uCount >= currentLimitRequest.count()) { // 超过次数，不执行目标方法
            //可以直接抛出异常
            return R.error("接口并发请求超过限制，请10秒后重试");
        } else if (uCount == 0){ // 第一次请求时，设置开始有效时间
            map.put(request.getRemoteAddr(), uCount + 1, ExpirationPolicy.CREATED, currentLimitRequest.time(), TimeUnit.MILLISECONDS);
        } else { // 未超过次数， 记录数据加一
            map.put(request.getRemoteAddr(), uCount + 1);
        }
        concurrentHashMap.put(request.getRequestURI(), map);

        // result的值就是被拦截方法的返回值
        return point.proceed();
    }

}

